# 起名大师，数仓格式为一个数组
import os

from dotenv import load_dotenv
from langchain import PromptTemplate
from langchain.llms import Tongyi
from langchain.schema import BaseOutputParser

load_dotenv("../ai.env")


# 自定义类
class CommaSeparatedListOutputParser(BaseOutputParser):
    """Parse the output of an LLM call to a comma-separated list."""

    def parse(self, text: str):
        """Parse the output of an LLM call."""
        print(text)
        return text.strip().split(",")


api_key = os.getenv("DASHSCOPE_API_KEY")
llm = Tongyi(
    model_name="qwen-max",
    temperature=0,
    dashscope_api_key=api_key
)
prompt = PromptTemplate.from_template(
    "你是一个起名大师，请模仿示例起3个{county}名字，比如男孩经常叫做{boy}，女孩经常叫做{girl}。请返回以逗号分隔的列表形式。仅返回逗号分隔的列表，不要返回其他内容。")
message = prompt.format(county="美国男孩", boy="sam", girl="lucy")
print(message)
response = llm.predict(message)
print(CommaSeparatedListOutputParser().parse(response))
